from heapq import *
from numpy import shape


# 合并集合
def union(i, j, A):
    i = i - 1
    j = j - 1
    # 使用加权规则合并根为i和j的两个集合i!=j
    x = A[i] + A[j] - 1
    if A[i] > A[j]:
        A[i] = j + 1
        A[j] = x
    else:
        A[j] = i + 1
        A[i] = x
    return A


# 寻找根节点：
def find(i, A):
    j = i - 1
    while A[j] > 0:
        j = A[j] - 1
    return j + 1


# 使用压缩规则的查找算法
def findUseCompress(i, A):
    j = i - 1
    while A[j] > 0:
        j = A[j] - 1
    k = i - 1
    while k != j:
        t = A[k]
        A[k] = j + 1
        k = t - 1
    return j + 1, A


# Kruskal算法
def kruskalArithmatic(cost):
    # 用来存放最小生成树
    T=[]
    # A中存放所有边的权重
    A=[]
    # A_nodes存放边对应的权重，边用两个结点表示
    A_nodes={}
    row, line=shape(cost)
    for r in range(row):
        for l in range(line):
            if r>l:
                A.append(cost[r, l])
                if cost[r, l] not in A_nodes:
                    A_nodes[cost[r, l]]=(r+1, l+1)
                else:
                    A_nodes[cost[r, l]+0.001] = (r + 1, l + 1)
    # B存放结点间的集合关系
    B = [0 for i in range(row)]
    # 小顶堆
    heapify(A)
    # 用来存结点
    nodes = [0 for i in range(row)]
    # 用来存放最小表的和
    mincost = 0
    i = 0
    while i < len(nodes) - 1 and len(A) != 0:
        minEdge = heappop(A)
        tupple=A_nodes[minEdge]
        l=find(tupple[0], B)
        k=find(tupple[1], B)
        if l!=k:
            i=i+1
            T.append(tupple)
            mincost=mincost+minEdge
            union(l, k, B)
    return T

